package net.xuele.learn.flink.example;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ReducingState;
import org.apache.flink.api.common.state.ReducingStateDescriptor;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.source.FromElementsFunction;

/**
 * @Author patrick
 * @Date 2023/7/12 18:49
 * @Description：
 *
 *  @see FromElementsFunction
 *
 */
public class CheckpointedFunctionTest<T> implements MapFunction<T, T>, CheckpointedFunction {

    private ReducingState<Long> countPerKey;

    private ListState<Long> countPerPartition;

    private long localCount;
    @Override
    public T map(T value) throws Exception {
        // update the states
        countPerKey.add(1L);
        localCount++;

        return value;
    }

    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        // the keyed state is always up to date anyways
        // just bring the per-partition state in shape
        countPerPartition.clear();
        countPerPartition.add(localCount);
    }

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {

        // get the state data structure for the per-key state
        countPerKey = context.getKeyedStateStore().getReducingState(
                new ReducingStateDescriptor<>("perKeyCount", new AddFunction<>(), Long.class));

        // get the state data structure for the per-partition state
        countPerPartition = context.getOperatorStateStore().getListState(
                new ListStateDescriptor<>("perPartitionCount", Long.class));

        // initialize the "local count variable" based on the operator state
        for (Long l : countPerPartition.get()) {
            localCount += l;
        }
    }

    private class AddFunction<Long> implements ReduceFunction<Long> {
        @Override
        public Long reduce(Long value1, Long value2) throws Exception {
            return null;
        }
    }
}
